# 研究机构：上海大学机自学院
# 研 究 生：王强
# 开发时间：2024/6/12 21:28
import numpy as np
from scipy.optimize import linear_sum_assignment
import math

class GNN:
    def __init__(self, targets, measurements):
        self.targets = targets
        self.measurements = measurements
        self.valid_measurements = []
        self.invalid_measurements = []
        self.association_matrix = []
        self.num_targets = 0                            # 目标的数量
        self.pd = 0.9                                   # 探测到目标的概率
        self.w = 2                                      # 测量维度，为2维
        self.beta = 0.0001                                # 新回波密度
        self.S = []                                     # 新息协方差
        self.threshold = 10                              # 关联门限

    def calculate_distance(self, target, measurement, S):
        # 计算目标和测量之间的欧氏距离
        a1 = target - measurement
        a2 = np.linalg.matrix_power(S, -1)  # np.linalg.matrix_power 表示为矩阵的求逆
        a3 = np.dot(a1, a2)
        d_ij = np.dot(a3, np.transpose(a1))
        return math.sqrt(d_ij)

    # 最近邻数据关联算法，先排除不在关联门限内的雷达目标量测
    def nearest_neighbor(self):
        """
        最近邻数据关联算法（利用跟踪门滤除无关目标）
        :param targets: 目标的坐标列表（Nx2矩阵，每行表示一个目标的坐标）
        :param measurements: 测量的坐标列表（Mx2矩阵，每行表示一个测量的坐标）
        :param threshold: 阈值，用于判断是否位于关联门限内
        :return: 门限内数组(有效量测)，不在门限内数组（无效量测），初步滤除无关目标，剩余无关目标用于航迹起始，判断舒服事新目标
        """
        self.num_targets = len(self.targets)
        num_measurements = len(self.measurements)
        valid_index = []
        for i in range(self.num_targets):
            for j in range(num_measurements):
                distance = self.calculate_distance(self.targets[i], self.measurements[j], self.S[i])
                b1 = (1 - self.pd) * pow(2*math.pi, self.w / 2) * self.beta * np.linalg.norm(self.S[i])
                b2 = self.pd / b1
                # print(b1)
                self.threshold = 2 * math.log(b2)
                if distance < self.threshold:
                    self.valid_measurements.append(self.measurements[j])
                    valid_index.append(j)
        # print(valid_index)
        self.invalid_measurements = [n for i, n in enumerate(self.measurements) if i not in valid_index]

    # 全局最近邻数据关联算法，获得量测与目标关联矩阵
    def global_nearest_neighbor(self, tracks):
        """
        全局最近邻数据关联算法
        :param targets: 目标的坐标列表（Nx2矩阵，每行表示一个目标的坐标）
        :param validation_measurements: 关联门限内测量的坐标列表（Mx2矩阵，每行表示一个测量的坐标）
        :return: 关联矩阵，每行表示一个目标与关联门限内测量的关联情况
        """
        # associated_track = {}
        self.nearest_neighbor()
        num_measurements = len(self.valid_measurements)
        cost_matrix = np.zeros((self.num_targets, num_measurements), dtype=float)
        self.association_matrix = np.zeros((self.num_targets, num_measurements), dtype=int)
        for t in range(self.num_targets):
            for m in range(num_measurements):
                cost_matrix[t, m] = np.linalg.norm(self.targets[t] - self.valid_measurements[m])
                # associated_track[t] = [self.targets[t]]
        # print(self.valid_measurements)
        row_ind, col_ind = linear_sum_assignment(cost_matrix)
        # print(row_ind)
        # print(col_ind)
        for i in range(len(row_ind)):
            self.association_matrix[row_ind[i], col_ind[i]] = 1
            # print(tracks[row_ind[i]]["points"])
            tracks[row_ind[i]]["associate_points"] += [self.valid_measurements[col_ind[i]]]
            # print(tracks[row_ind[i]]["points"])
            # associated_track[row_ind[i]].append(self.valid_measurements[col_ind[i]])
            # print(self.valid_measurements)
        return tracks

if __name__ == '__main__':

    # 示例数据
    targets = np.array([[1.0, 2.0], [3.0, 4.0], [2.0, 6.0], [10, 10], [100, 10]])
    measurements = np.array([[3.2, 4.3], [10, 11], [1.5, 2.5], [5.5, 6.9], [15.5, 6.8], [5.4, 6.8]])
    tracks = []
    for t in range(len(targets)):
        tracks.append({
            'id': t,
            'points': [targets[t]]
        })
    print(tracks)
    # 调用全局最近邻算法
    # 第一次滤除无关杂波
    GNN_mtt = GNN(targets=targets, measurements=measurements)
    tracks = GNN_mtt.global_nearest_neighbor(tracks)
    # print(tracks)
    # col_sums = np.sum(association_matrix, axis=0)
    # zero_indices = np.where(col_sums == 0)[0]
    # print(zero_indices)

    # for i in zero_indices:
    #     GNN_mtt.invalid_measurements.append(GNN_mtt.valid_measurements[i])
    # valid_measurements = [n for i, n in enumerate(GNN_mtt.valid_measurements) if i not in zero_indices]
    # print(valid_measurements)
    # print(invalid_measurements)
    # row_ind, col_ind = linear_sum_assignment(cost_matrix)

    # print(row_ind, col_ind)
    # # 输出匹配结果
    # print("匹配结果:")
    # for i in range(len(row_ind)):
    #     print(f"目标 {row_ind[i]} 匹配到测量 {col_ind[i]} (距离: {cost_matrix[row_ind[i], col_ind[i]]:.2f})")
    #
    # # 显示未匹配的测量
    # unmatched_measurements = set(range(len(measurements))) - set(col_ind)
    # if unmatched_measurements:
    #     print("\n未匹配的测量点:", unmatched_measurements)
    # else:
    #     print("\n所有测量点都已匹配。")
